# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: inkyu

import datetime
import logging

import tabulate

from coin.proto.coin_query_pb2 import CurrencyBalance


class SyntheticFxProduct(object):
  def __init__(self, base, native_products):
    self.base = base
    self.native_products = native_products
    self.symbol = self.native_products[0].symbol
    for native_product in native_products:
      assert base.currency == native_product.base.currency, (base, native_product.base)


class SyntheticFxBalance(object):
  def __init__(self, *, base_balances, quote_balances, base_reserve):
    if quote_balances is not None:
      assert len(base_balances) == len(quote_balances)
    base_total = sum(base_balances)
    base_offset = base_total - base_reserve
    self.base_balances = base_balances
    self.quote_balances = quote_balances
    self.base_total = base_total
    self.base_reserve = base_reserve
    self.base_offset = base_offset


class SyntheticFxInfoPrinter(object):
  def __init__(self, ogs, quote_currencies, base_currencies):
    self._logger = logging.getLogger(__name__)

    self._ogs = ogs
    self._quote_currencies = quote_currencies
    assert len(self._ogs) == len(self._quote_currencies)
    self._base_currencies = base_currencies
    self._ogs_name = [og._og_info.CurrencyType._exchange for og in ogs]

  def print_balance(self):
    table = []
    for i, cur in enumerate(self._quote_currencies):
      bals = [og.get_balance(cur.currency) if i == j else 0 for j, og in enumerate(self._ogs)]
      table.append(['*', cur.currency, sum(bals)] + bals)
    for cur in self._base_currencies:
      # use cur.currency, not cur, as ogs do not have same type
      bals = [og.get_balance(cur.currency) for og in self._ogs]
      table.append(['*', cur.currency, sum(bals)] + bals)

    self._logger.info(
        '----- Balance:\n* %s\n%s' %
        (str(datetime.datetime.now()),
         tabulate.tabulate(
             table, headers=['*', 'Currency', 'Total'] + self._ogs_name, floatfmt=".6f")))

  def get_balance(self, arb_product, reserve_balance):
    assert type(arb_product) == SyntheticFxProduct
    # use cur.currency, not cur, as ogs do not have same type
    base_bals = [og.get_balance(arb_product.base.currency) for og in self._ogs]
    quote_bals = []
    for og, quote_currency in zip(self._ogs, self._quote_currencies):
      quote_bals.append(og.get_balance(quote_currency.currency))
    for base_bal in base_bals:
      if base_bal is None:
        return None
    for quote_bal in quote_bals:
      if quote_bal is None:
        return None
    return SyntheticFxBalance(base_balances=base_bals,
                              quote_balances=quote_bals,
                              base_reserve=reserve_balance)


class KimchiEstimatedBalancePrinter(object):
  def __init__(self,
               *,
               quote_currencies,
               base_currencies,
               extra_currencies=None,
               base_currency_reserves,
               ogs,
               use_exchange1=False):
    self._check_init_param_consistency(quote_currencies, base_currencies, base_currency_reserves)
    self._quote_currencies = quote_currencies
    self._extra_currencies = extra_currencies
    self._base_currencies = base_currencies
    self._base_currency_reserves = base_currency_reserves
    self._ogs = ogs
    self._use_exchange1 = use_exchange1
    self._pos_amts = None

    self._balance_printer = SyntheticFxInfoPrinter(self._ogs,
                                                   self._quote_currencies,
                                                   self._base_currencies)
    self._product_prices = {}
    self._logger = logging.getLogger(__name__)

  def calculate_pnl_balance(self):
    if not self._ogs:
      return None
    quote0_bal = 0.0
    quote1_bal = 0.0
    for base_currency in self._base_currencies:
      price = self._product_prices.get(base_currency, None)
      if price is None:
        return None
      reserve = self._base_currency_reserves[base_currency]
      balance = sum([og.get_balance(base_currency) for og in self._ogs])
      if self._use_exchange1:
        quote1_bal += (balance - reserve) * price
      else:
        quote0_bal += (balance - reserve) * price
    quote0_bal += self._ogs[0].get_balance(self._quote_currencies[0])
    quote1_bal += self._ogs[1].get_balance(self._quote_currencies[1])
    pnl_balance0 = CurrencyBalance(currency=self._quote_currencies[0].currency, total=quote0_bal)
    pnl_balance1 = CurrencyBalance(currency=self._quote_currencies[1].currency, total=quote1_bal)
    return [pnl_balance0, pnl_balance1]

  @staticmethod
  def _check_init_param_consistency(quote_currencies, base_currencies, base_currency_reserves):
    assert set(base_currencies) == set(base_currency_reserves)

  def get_balance(self, arb_product, reserve_balance):
    return self._balance_printer.get_balance(arb_product, reserve_balance)

  def on_book(self, product, exchange_id, book):
    if not self._use_exchange1 and exchange_id == 0:
      # Only use exchange 0's price for estimation.
      assert product.quote == self._quote_currencies[0]
      mid_price = 0.5 * (book.ask0().price + book.bid0().price)

      currency = product.base
      self._product_prices[currency] = mid_price
    elif self._use_exchange1 and exchange_id == 1:
      # Only use exchange 1's price for estimation.
      assert product.quote == self._quote_currencies[0]
      mid_price = 0.5 * (book.ask0().price + book.bid0().price)

      currency = product.base
      self._product_prices[currency] = mid_price

  def set_pos_amts_total(self, pos_amts):
    self._pos_amts = pos_amts

  def has_pos_amts(self):
    return self._pos_amts is not None

  def print_balance(self, implied_price):
    table = []
    for i, cur in enumerate(self._quote_currencies):
      balances = [og.get_balance(cur) if i == j else 0 for j, og in enumerate(self._ogs)]
      for balance in balances:
        if balance is None:
          return
      table.append({
          'currency': cur,
          'total': sum(balances),
          **{og.exchange: balances[i] for i, og in enumerate(self._ogs)}
      })

    if self._extra_currencies:
      for i, extra_currencies in self._extra_currencies.items():
        for cur, value in extra_currencies.items():
          balances = [og.get_balance(cur) if i == j else 0 for j, og in enumerate(self._ogs)]
          for balance in balances:
            if balance is None:
              return
          table.append({
              'currency': cur,
              'total': sum(balances),
              **{og.exchange: balances[i] for i, og in enumerate(self._ogs)}
          })

    for cur in self._base_currencies:
      balances = [og.get_balance(cur) for og in self._ogs]
      for balance in balances:
        if balance is None:
          return
      table.append({
          'currency': cur,
          'total': sum(balances),
          **{og.exchange: balances[i] for i, og in enumerate(self._ogs)}
      })

    okdk = sum([og.get_balances().get('OKDK', 0.0) for og in self._ogs])
    bnb = sum([og.get_balances().get('BNB', 0.0) for og in self._ogs])

    quote = 0
    adjusted_quote = 0
    estimated_base = 0
    adjusted_estimated_base = 0
    imbalance_total = 0.
    extra_value_1 = 0.
    for entry in table:
      currency = entry['currency']
      reserve = self._base_currency_reserves.get(currency, 0.)
      price = self._product_prices.get(currency, 0.)
      if self._extra_currencies and entry['currency'] in self._extra_currencies.get(1, {}):
        # Only support aggressive
        price = self._extra_currencies[1][entry['currency']] * self._product_prices.get(
            self._quote_currencies[1])
        entry['reserve'] = reserve
        entry['imbalance'] = 0
        entry['imbalance_amount'] = 0
        entry['estimated_base'] = 0
        entry['adjusted_estimated_base'] = 0
        adjusted_quote += entry['total'] * price
        extra_value_1 += entry['total'] * price
      elif entry['currency'] not in self._quote_currencies:
        entry['reserve'] = reserve
        entry['imbalance'] = entry['total'] - entry['reserve']
        entry['imbalance_amount'] = entry['imbalance'] * price
        entry['estimated_base'] = entry['total'] * price
        entry['adjusted_estimated_base'] = entry['reserve'] * price
        adjusted_quote += entry['imbalance_amount']
        imbalance_total += entry['imbalance_amount']
        estimated_base += entry['estimated_base']
        adjusted_estimated_base += entry['adjusted_estimated_base']
      else:
        entry['reserve'] = reserve
        entry['imbalance'] = 0
        entry['imbalance_amount'] = 0
        entry['estimated_base'] = 0
        entry['adjusted_estimated_base'] = 0
        adjusted_quote += entry['total'] * price
        quote += entry['total'] * price

    quote0 = self._ogs[0].get_balance(self._quote_currencies[0])
    if len(self._ogs) > 1:
      quote1 = self._ogs[1].get_balance(self._quote_currencies[1]) + extra_value_1
    else:
      quote1 = None

    headers = (['Currency'] + [og.exchange for og in self._ogs] + [
        'Total',
        'Reserve',
        'Imbalance',
        'ImbalanceAmount',
        'EstimatedBase',
        'AdjustedEstimatedBase'
    ])
    fmt_table = []
    for entry in table:
      fmt_table.append([entry['currency']] + [entry[og.exchange] for og in self._ogs] + [
          entry['total'],
          entry['reserve'],
          entry['imbalance'],
          entry['imbalance_amount'],
          entry['estimated_base'],
          entry['adjusted_estimated_base']
      ])

    if quote1:
      res = (
          "%s\n %s: %.4f, %s: %.4f, Imbalance: %.4f, Adjusted%s: %.4f, EstimatedBase: %.4f, OKDK: %.4f, BNB: %.4f\n"
          "Position%s: %s (%.2f%%), Position%s: %s (%.2f%%) \n"
          "Balance%s: %s Balance%s: %s \n"
          "Reserve%s: %s Reserve%s: %s \n"
          "Total %s: %s \n") % (tabulate.tabulate(
              fmt_table, headers, floatfmt='.4f', stralign="right"),
                                self._quote_currencies[1].currency,
                                quote1,
                                self._quote_currencies[0].currency,
                                quote0,
                                imbalance_total,
                                self._quote_currencies[0].currency,
                                quote0 + imbalance_total,
                                estimated_base,
                                okdk,
                                bnb,
                                self._quote_currencies[0].currency,
                                self._pos_amts.position_amt_0 if self._pos_amts else 0.,
                                (100. * self._pos_amts.position_amt_0
                                 / self._pos_amts.reserve_amt_0 if self._pos_amts else 0.),
                                self._quote_currencies[1].currency,
                                self._pos_amts.position_amt_1 if self._pos_amts else 0.,
                                (100. * self._pos_amts.position_amt_1
                                 / self._pos_amts.reserve_amt_1 if self._pos_amts else 0.),
                                self._quote_currencies[0].currency,
                                self._pos_amts.balance_amt_0 if self._pos_amts else 0.,
                                self._quote_currencies[1].currency,
                                self._pos_amts.balance_amt_1 if self._pos_amts else 0.,
                                self._quote_currencies[0].currency,
                                self._pos_amts.reserve_amt_0 if self._pos_amts else 0.,
                                self._quote_currencies[1].currency,
                                self._pos_amts.reserve_amt_1 if self._pos_amts else 0.,
                                self._quote_currencies[0].currency,
                                quote0 + imbalance_total + quote1 * implied_price)
    else:
      res = (
          "%s\n%s: %.4f, Imbalance: %.4f, Adjusted%s: %.4f, EstimatedBase: %.4f, OKDK: %.4f, BNB: %.4f\n"
          "Position%s: %s (%.2f%%) \n"
          "Balance%s: %s \n"
          "Reserve%s: %s \n") % (
              tabulate.tabulate(fmt_table, headers, floatfmt='.4f', stralign="right"),
              self._quote_currencies[0].currency,
              quote0,
              imbalance_total,
              self._quote_currencies[0].currency,
              quote0 + imbalance_total,
              estimated_base,
              okdk,
              bnb,
              self._quote_currencies[0].currency,
              self._pos_amts.position_amt_0 if self._pos_amts else 0.,
              (100. * self._pos_amts.position_amt_0
               / self._pos_amts.reserve_amt_0 if self._pos_amts else 0.),
              self._quote_currencies[0].currency,
              self._pos_amts.balance_amt_0 if self._pos_amts else 0.,
              self._quote_currencies[0].currency,
              self._pos_amts.reserve_amt_0 if self._pos_amts else 0.,
          )
    self._logger.info('----- Estimated Balance:\n* %s\n%s' % (str(datetime.datetime.now()), res))
    return imbalance_total, quote0, quote1
